import { SendMessageProps, UniversalChunkProps } from "../types";
import { BaseProvider } from "./BaseProvider";
import OpenAI from "openai";

export class QianfanProvider extends BaseProvider {
  private client: OpenAI;
  constructor(apiKey: string, baseURL: string) {
    super();
    this.client = new OpenAI({
      apiKey: apiKey, // process.env.OPENAI_API_KEY, // process.env.ALI_API_KEY,
      baseURL: baseURL, //"https://qianfan.baidubce.com/v2/", //"https://dashscope.aliyuncs.com/compatible-mode/v1", // "https://qianfan.baidubce.com/v2/", // 千帆ModelBuilder平台地址
    });
  }
  async chat(messages: SendMessageProps[], model: string) {
    const stream = await this.client.chat.completions.create({
      messages: messages as any, //[{ role: "user", content }],
      model,
      stream: true,
    });
    const self = this;
    return {
      async *[Symbol.asyncIterator]() {
        for await (const chunk of stream) {
          yield self.transformResponse(chunk);
        }
      },
    };
  }
  protected transformResponse(chunk: any): UniversalChunkProps {
    return {
      is_end: chunk.choices[0].finish_reason === "stop",
      result: chunk.choices[0].delta.content || "",
    };
  }
}
